#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
多激光雷达 LaserScan -> PointCloud2（统一到 target_frame，默认 base_link）-> 合并为单一话题
稳健性增强：
- 事件驱动：任一雷达来一帧就尝试合并发布（受 max_publish_rate_hz 限制）
- 缺数据/TF 不可用：静默跳过该路，不警告
- TF 默认取最新（Time(0)），适配静态外参，避免 rosbag 的过去/将来外推报错
- 支持过期剔除（max_staleness_sec）
- 可选打印各 /scan 实得频率（节流输出）
"""

import os
import yaml
import math
import rospy
from sensor_msgs.msg import LaserScan, PointCloud2, PointField
from sensor_msgs import point_cloud2 as pc2
from laser_geometry import LaserProjection

import tf2_ros
from tf2_sensor_msgs.tf2_sensor_msgs import do_transform_cloud


def has_field(cloud, name: str) -> bool:
    return any(f.name == name for f in cloud.fields)


class RateMeter:
    """按 header.stamp 估计频率（EMA），并按 wall-clock 节流打印"""
    def __init__(self, topic, alpha=0.2, print_every_sec=2.0, enable=True):
        self.topic = topic
        self.alpha = alpha
        self.enable = enable
        self.last_stamp = None
        self.period_ema = None
        self.last_print_wall = rospy.get_time()
        self.print_every = float(print_every_sec)

    def update(self, stamp):
        if not self.enable:
            return
        t = stamp.to_sec()
        if self.last_stamp is None:
            self.last_stamp = t
            return
        dt = t - self.last_stamp
        self.last_stamp = t
        if dt <= 0 or not math.isfinite(dt):
            return
        if self.period_ema is None:
            self.period_ema = dt
        else:
            self.period_ema = self.alpha * dt + (1 - self.alpha) * self.period_ema

        now = rospy.get_time()
        if now - self.last_print_wall >= self.print_every:
            hz = 1.0 / self.period_ema if self.period_ema and self.period_ema > 0 else 0.0
            rospy.loginfo_throttle(0.5, "[%s] ~%.2f Hz", self.topic, hz)
            self.last_print_wall = now


class SingleLaserProcessor(object):
    """单路雷达：订阅 Scan，投影->变换到 target_frame，存最新帧；可选打印频率"""

    def __init__(self, scan_topic, target_frame, tf_buffer,
                 tf_timeout_sec, use_latest_tf, fallback_to_latest_on_fail,
                 store_fn, print_rate, print_rate_every_sec):
        self.scan_topic = scan_topic
        self.target_frame = target_frame
        self.tf_buffer = tf_buffer
        self.tf_timeout = rospy.Duration.from_sec(float(tf_timeout_sec))
        self.use_latest_tf = bool(use_latest_tf)
        self.fallback_to_latest_on_fail = bool(fallback_to_latest_on_fail)

        self.store_fn = store_fn
        self.rate_meter = RateMeter(scan_topic, enable=print_rate, print_every_sec=print_rate_every_sec)

        self.laser_projection = LaserProjection()
        self.sub = rospy.Subscriber(self.scan_topic, LaserScan, self.cb, queue_size=10)

        rospy.loginfo("订阅: %s  (target_frame=%s, TF=%s)",
                      self.scan_topic, self.target_frame,
                      "latest" if self.use_latest_tf else "stamp")

    def _lookup_tf(self, source_frame, stamp):
        # 首选策略
        query_time = rospy.Time(0) if self.use_latest_tf else stamp
        if self.tf_buffer.can_transform(self.target_frame, source_frame, query_time, rospy.Duration(0.0)):
            return self.tf_buffer.lookup_transform(self.target_frame, source_frame, query_time, self.tf_timeout)

        # 可选回退：失败时再用 latest
        if (not self.use_latest_tf) and self.fallback_to_latest_on_fail:
            qt = rospy.Time(0)
            if self.tf_buffer.can_transform(self.target_frame, source_frame, qt, rospy.Duration(0.0)):
                return self.tf_buffer.lookup_transform(self.target_frame, source_frame, qt, self.tf_timeout)
        raise RuntimeError("TF unavailable")

    def cb(self, scan_msg: LaserScan):
        try:
            self.rate_meter.update(scan_msg.header.stamp)

            # 1) Scan -> Cloud（源帧）
            cloud_local = self.laser_projection.projectLaser(scan_msg)

            # 2) 源帧即目标帧：直接存
            if self.target_frame == cloud_local.header.frame_id:
                cloud_target = cloud_local
            else:
                # 3) 查 TF（静默失败）
                try:
                    tfm = self._lookup_tf(cloud_local.header.frame_id, scan_msg.header.stamp)
                except Exception:
                    return  # 静默跳过该帧
                cloud_target = do_transform_cloud(cloud_local, tfm)

            # 4) 存最新帧
            self.store_fn(self.scan_topic, cloud_target)

        except Exception:
            # 静默：不打印 warn/err，保证“有没有都一样”
            return


class MultiLaserMergedNode(object):
    """管理多路雷达 + 事件驱动合并输出（静默容错）"""

    def __init__(self):
        rospy.init_node('multi_laser_points_merged', anonymous=True)

        # ---------- 配置 ----------
        self.config = self.load_config()
        self.target_frame = self.config.get('target_frame', 'base_link')
        self.merge_output_topic = self.config.get('merge_output_topic', '/laser/points_merged')

        # 发布控制：事件驱动 + 上限速率
        self.max_publish_rate_hz = float(self.config.get('max_publish_rate_hz', 30))
        self.min_pub_interval = 1.0 / self.max_publish_rate_hz if self.max_publish_rate_hz > 0 else 0.0
        self.last_pub_wall = 0.0

        # 数据选择
        self.include_intensity = bool(self.config.get('include_intensity', True))
        self.max_staleness_sec = float(self.config.get('max_staleness_sec', 0.25))  # 超过则视为过期不参与合并

        # TF 策略
        self.tf_timeout_sec = self.config.get('tf_timeout_sec', 0.05)
        self.use_latest_tf = bool(self.config.get('use_latest_tf', True))  # 默认 true：最稳
        self.fallback_to_latest_on_fail = bool(self.config.get('fallback_to_latest_on_fail', True))

        # 打印频率
        self.print_scan_rate = bool(self.config.get('print_scan_rate', True))
        self.print_scan_rate_every_sec = float(self.config.get('print_scan_rate_every_sec', 2.0))

        # 话题列表
        self.scan_topics = self.parse_scan_topics(self.config)

        # ---------- 基础设施 ----------
        self.tf_buffer = tf2_ros.Buffer(cache_time=rospy.Duration(10.0))
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer)

        self.latest_clouds = {}      # topic -> PointCloud2（在 target_frame 下）
        self.latest_walltime = {}    # topic -> wall-clock 到达时间（用于过期判断）
        self.pub = rospy.Publisher(self.merge_output_topic, PointCloud2, queue_size=1)

        # 创建处理器（订阅）
        self.processors = [
            SingleLaserProcessor(
                scan_topic=t,
                target_frame=self.target_frame,
                tf_buffer=self.tf_buffer,
                tf_timeout_sec=self.tf_timeout_sec,
                use_latest_tf=self.use_latest_tf,
                fallback_to_latest_on_fail=self.fallback_to_latest_on_fail,
                store_fn=self.store_latest_and_try_publish,
                print_rate=self.print_scan_rate,
                print_rate_every_sec=self.print_scan_rate_every_sec
            ) for t in self.scan_topics
        ]

        rospy.loginfo("合并 %d 路 → %s  (frame=%s, event-driven, max %.1f Hz, staleness=%.2fs)",
                      len(self.scan_topics), self.merge_output_topic,
                      self.target_frame, self.max_publish_rate_hz, self.max_staleness_sec)

    # ----- 配置 -----
    def load_config(self):
        cfg_file = rospy.get_param('~config_file', 'config/laser_topics_config.yaml')
        package_name = rospy.get_param('~package_name', 'savior_switcher')

        if not os.path.isabs(cfg_file):
            try:
                import rospkg
                package_path = rospkg.RosPack().get_path(package_name)
                cfg_file = os.path.join(package_path, cfg_file)
            except Exception:
                pass

        try:
            with open(cfg_file, 'r') as f:
                cfg = yaml.safe_load(f) or {}
            rospy.loginfo("加载配置: %s", cfg_file)
            return cfg
        except Exception as e:
            rospy.logwarn("加载配置失败，将使用最小默认: %s", str(e))
            return {
                'target_frame': 'base_link',
                'merge_output_topic': '/laser/points_merged',
                'laser_inputs': ['/scan', '/scan2']
            }

    @staticmethod
    def parse_scan_topics(cfg_dict):
        topics = []
        if 'laser_inputs' in cfg_dict:
            topics = [t for t in cfg_dict.get('laser_inputs', []) if isinstance(t, str) and t]
        elif 'laser_conversions' in cfg_dict:
            for item in cfg_dict.get('laser_conversions', []):
                if isinstance(item, dict) and item.get('scan_topic'):
                    topics.append(item['scan_topic'])
        # 去重
        dedup = []
        for t in topics:
            if t not in dedup:
                dedup.append(t)
        return dedup or ['/scan']

    # ----- 数据到达 → 尝试发布 -----
    def store_latest_and_try_publish(self, topic, cloud_in_target):
        self.latest_clouds[topic] = cloud_in_target
        self.latest_walltime[topic] = rospy.get_time()
        self.try_publish()

    def try_publish(self):
        # 速率限制
        if self.min_pub_interval > 0:
            now = rospy.get_time()
            if now - self.last_pub_wall < self.min_pub_interval:
                return

        # 选取“不过期”的最新帧
        now = rospy.get_time()
        valid = []
        for t, c in self.latest_clouds.items():
            ts = self.latest_walltime.get(t, 0)
            if (now - ts) <= self.max_staleness_sec:
                valid.append(c)

        if not valid:
            return  # 所有路都过期/无数据：不发布

        # 字段：xyz 必有；如任一路有 intensity 且允许，则带上 intensity（无的填 0）
        any_has_intensity = any(has_field(c, 'intensity') for c in valid)
        with_intensity = self.include_intensity and any_has_intensity

        fields = [
            PointField('x', 0,  PointField.FLOAT32, 1),
            PointField('y', 4,  PointField.FLOAT32, 1),
            PointField('z', 8,  PointField.FLOAT32, 1),
        ]
        if with_intensity:
            fields.append(PointField('intensity', 12, PointField.FLOAT32, 1))

        # 时间戳：取参与合并帧中“最新”的 stamp；frame 统一为 target_frame
        latest_stamp = max(c.header.stamp for c in valid)
        header = valid[0].header
        header.stamp = latest_stamp
        header.frame_id = self.target_frame

        # 聚合
        merged = []
        if with_intensity:
            for c in valid:
                if has_field(c, 'intensity'):
                    for x, y, z, inten in pc2.read_points(c, field_names=['x', 'y', 'z', 'intensity'], skip_nans=True):
                        merged.append((x, y, z, float(inten)))
                else:
                    for x, y, z in pc2.read_points(c, field_names=['x', 'y', 'z'], skip_nans=True):
                        merged.append((x, y, z, 0.0))
        else:
            for c in valid:
                for x, y, z in pc2.read_points(c, field_names=['x', 'y', 'z'], skip_nans=True):
                    merged.append((x, y, z))

        if not merged:
            return

        cloud = pc2.create_cloud(header, fields, merged)
        self.pub.publish(cloud)
        self.last_pub_wall = rospy.get_time()

    def run(self):
        rospy.spin()


if __name__ == '__main__':
    try:
        MultiLaserMergedNode().run()
    except rospy.ROSInterruptException:
        pass
